"""
大模型回答，生成json格式数据
"""
import re
import openai
from core.config import settings
from typing import List, Dict, Union
import json
import json_repair


async def answer_with_llm(prompt: str, parse_type: str = 'dict')-> Union[Dict, List[Dict]]:

    openai_client = openai.AsyncClient(
        base_url=settings.LLM_URI,
        api_key=settings.LLM_API_KEY,
        default_headers=json.loads(settings.LLM_HEADERS)
    )
    messages = [
        {'role': 'system', 'content': '你是一个合同审查员'},
        {'role': 'user', 'content': prompt}
    ]
    openai_args = {
        'model': settings.LLM_MODEL,
        'messages': messages,
        'temperature': 0,
        'top_p': 0.1,
        'max_tokens': None,
        'stop': None
    }
    response = await openai_client.chat.completions.create(**openai_args)
    response_content = response.choices[0].message.content

    if parse_type == 'dict':
        result = parse_result_dict(response_content)
    elif parse_type == 'list':
        result = parse_result_list(response_content)
    else:
        result = parse_result_str(response_content)
    return result


def parse_result_dict(text: str)-> Dict:
    result = {}
    pattern = r"```json(.*?)```"
    matches = re.findall(pattern, text, re.DOTALL)
    for match in matches:
        fixed_json = json_repair.repair_json(match, ensure_ascii=False, return_objects=True)
        if isinstance(fixed_json, dict):
            result.update(fixed_json)
    return result


def parse_result_list(text: str) -> List[Dict]:
    result = []
    pattern = r"```json(.*?)```"
    matches = re.findall(pattern, text, re.DOTALL)
    for match in matches:
        fixed_json = json_repair.repair_json(match, ensure_ascii=False, return_objects=True)
        if isinstance(fixed_json, dict):
            result.append(fixed_json)
        if isinstance(fixed_json, list):
            result.extend(fixed_json)
    return result


def parse_result_str(text: str) -> str:
    """去除<think></think>部分的内容"""
    return re.split('</think>', text, re.MULTILINE)[-1].strip()
